#include <iostream>
#include <cstring>
#include <chrono>
#include <glm/mat4x4.hpp>
#include <glm/gtx/transform.hpp>
#include "VK_UniformBuffer.h"
#include "VK_Context.h"
#include "VK_Pipeline.h"
#include "VK_Image.h"
#include "VK_Texture.h"
#include "VK_DynamicState.h"
#include "VK_Window.h"
#include "VK_CommandPool.h"
#include "../Node.h"
#include "../ObjectFactory.h"

using namespace std;

const std::vector<float> vertices = {
    -0.5f, -0.5f, -0.5f,  0.0f, 0.0f,
        0.5f, -0.5f, -0.5f,  1.0f, 0.0f,
        0.5f,  0.5f, -0.5f,  1.0f, 1.0f,
        0.5f,  0.5f, -0.5f,  1.0f, 1.0f,
        -0.5f,  0.5f, -0.5f,  0.0f, 1.0f,
        -0.5f, -0.5f, -0.5f,  0.0f, 0.0f,

        -0.5f, -0.5f,  0.5f,  0.0f, 0.0f,
        0.5f, -0.5f,  0.5f,  1.0f, 0.0f,
        0.5f,  0.5f,  0.5f,  1.0f, 1.0f,
        0.5f,  0.5f,  0.5f,  1.0f, 1.0f,
        -0.5f,  0.5f,  0.5f,  0.0f, 1.0f,
        -0.5f, -0.5f,  0.5f,  0.0f, 0.0f,

        -0.5f,  0.5f,  0.5f,  1.0f, 0.0f,
        -0.5f,  0.5f, -0.5f,  1.0f, 1.0f,
        -0.5f, -0.5f, -0.5f,  0.0f, 1.0f,
        -0.5f, -0.5f, -0.5f,  0.0f, 1.0f,
        -0.5f, -0.5f,  0.5f,  0.0f, 0.0f,
        -0.5f,  0.5f,  0.5f,  1.0f, 0.0f,

        0.5f,  0.5f,  0.5f,  1.0f, 0.0f,
        0.5f,  0.5f, -0.5f,  1.0f, 1.0f,
        0.5f, -0.5f, -0.5f,  0.0f, 1.0f,
        0.5f, -0.5f, -0.5f,  0.0f, 1.0f,
        0.5f, -0.5f,  0.5f,  0.0f, 0.0f,
        0.5f,  0.5f,  0.5f,  1.0f, 0.0f,

        -0.5f, -0.5f, -0.5f,  0.0f, 1.0f,
        0.5f, -0.5f, -0.5f,  1.0f, 1.0f,
        0.5f, -0.5f,  0.5f,  1.0f, 0.0f,
        0.5f, -0.5f,  0.5f,  1.0f, 0.0f,
        -0.5f, -0.5f,  0.5f,  0.0f, 0.0f,
        -0.5f, -0.5f, -0.5f,  0.0f, 1.0f,

        -0.5f,  0.5f, -0.5f,  0.0f, 1.0f,
        0.5f,  0.5f, -0.5f,  1.0f, 1.0f,
        0.5f,  0.5f,  0.5f,  1.0f, 0.0f,
        0.5f,  0.5f,  0.5f,  1.0f, 0.0f,
        -0.5f,  0.5f,  0.5f,  0.0f, 0.0f,
        -0.5f,  0.5f, -0.5f,  0.0f, 1.0f
    };

class cube : public Node
{
public:
    cube()
    {
        nodeName = "cube";
    };
public:
    RunningState run()override
    {
        auto window = VK_Window::getInstance()->getWindow();
        auto title = getNodeTitle();
        glfwSetWindowTitle(window, title.data());

        context = createVkContext();

        VK_Context::VK_Config vkConfig;
        vkConfig.usingDepthAttachment = true;
        bool ok = context->initVulkanDevice(vkConfig);
        if (!ok)
            return RunningState_Failed;

        context->setClearColor(0.21, 0.21, 0.5, 1);

        auto shaderSet = context->createShaderSet();
        shaderSet->addShader("../shader/cube/vert.spv", VK_SHADER_STAGE_VERTEX_BIT);
        shaderSet->addShader("../shader/cube/frag.spv", VK_SHADER_STAGE_FRAGMENT_BIT);

        shaderSet->appendVertexAttributeDescription(0, sizeof (float) * 3, VK_FORMAT_R32G32B32_SFLOAT, 0);
        shaderSet->appendVertexAttributeDescription(1, sizeof (float) * 2, VK_FORMAT_R32G32_SFLOAT,
                                                    sizeof(float) * 3);

        shaderSet->appendVertexInputBindingDescription(5 * sizeof(float), 0, VK_VERTEX_INPUT_RATE_VERTEX);

        VkDescriptorSetLayoutBinding uniformBinding = VK_ShaderSet::createDescriptorSetLayoutBinding(0,
                                                                                                     VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, VK_SHADER_STAGE_VERTEX_BIT);
        shaderSet->addDescriptorSetLayoutBinding(uniformBinding);

        auto samplerBinding = VK_ShaderSet::createDescriptorSetLayoutBinding(1,
                                                                             VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, VK_SHADER_STAGE_FRAGMENT_BIT);
        auto samplerCreateInfo  = VK_Sampler::createSamplerCreateInfo();
        auto samplerPtr = context->createSampler(samplerCreateInfo);
        VkSampler sampler = samplerPtr->getSampler();
        samplerBinding.pImmutableSamplers = &sampler;

        shaderSet->addDescriptorSetLayoutBinding(samplerBinding);

        if (!shaderSet->isValid()) {
            std::cerr << "invalid shaderSet" << std::endl;
            shaderSet->release();
            context->release();
            return RunningState_Failed;
        }

        auto ubo = shaderSet->addUniformBuffer(0, sizeof(float) * 16);
        ubo->setWriteDataCallback([this](char *&data, uint32_t size)->uint32_t{return updateUniformBufferData(data, size);});

        auto image = context->createImage("../images/smile.png", false);
        auto background = context->createImage("../images/container.jpg", false);

        auto command = context->getCommandPool()->beginSingleTimeCommands();

        adjustImageLayout(command, image->getImage(), VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
                          VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL);
        adjustImageLayout(command, background->getImage(), VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
                          VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);

        VkImageBlit region;
        region.srcSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
        region.srcSubresource.mipLevel = 0;
        region.srcSubresource.baseArrayLayer = 0;
        region.srcSubresource.layerCount = 1;
        region.srcOffsets[0].x = 0;
        region.srcOffsets[0].y = 0;
        region.srcOffsets[0].z = 0;
        region.srcOffsets[1].x = 512;
        region.srcOffsets[1].y = 512;
        region.srcOffsets[1].z = 1;
        region.dstSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
        region.dstSubresource.mipLevel = 0;
        region.dstSubresource.baseArrayLayer = 0;
        region.dstSubresource.layerCount = 1;
        region.dstOffsets[0].x = 0;
        region.dstOffsets[0].y = 0;
        region.dstOffsets[0].z = 0;
        region.dstOffsets[1].x = 256;
        region.dstOffsets[1].y = 256;
        region.dstOffsets[1].z = 1;

        vkCmdBlitImage(command, image->getImage(), VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,
                       background->getImage(),
                       VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
                       1, &region, VK_FILTER_LINEAR);

        adjustImageLayout(command, background->getImage(), VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
                          VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);

        context->getCommandPool()->endSingleTimeCommands(command, context->getGraphicQueue());


        auto imageViewCreateInfo = VK_ImageView::createImageViewCreateInfo(background->getImage(),
                                                                           VK_FORMAT_R8G8B8A8_SRGB);
        auto imageView = context->createImageView(imageViewCreateInfo);
        shaderSet->addImageView(imageView);

        context->initVulkanContext();

        pipeline = context->createPipeline(shaderSet);

        auto rasterCreateInfo = pipeline->getRasterizationStateCreateInfo();
        rasterCreateInfo.cullMode = VK_CULL_MODE_NONE;
        pipeline->setRasterizationStateCreateInfo(rasterCreateInfo);

        pipeline->create();
        auto buffer = context->createVertexBuffer(vertices, 5);
        pipeline->addRenderBuffer(buffer);

        context->createCommandBuffers();

        auto result = context->run();
        setAverageFps(context->getFps());
        context->release();
        return result;
    }

    uint32_t updateUniformBufferData(char *&data, uint32_t size)
    {
        static auto startTime = std::chrono::high_resolution_clock::now();
        auto currentTime = std::chrono::high_resolution_clock::now();
        float time = std::chrono::duration<float, std::chrono::seconds::period>
                     (currentTime - startTime).count();
        glm::mat4 model = glm::rotate(glm::mat4(1.0f), time * glm::radians(30.0f), glm::vec3(0.0f, 0.0f,
                                                                                             1.0f));

        auto view = glm::lookAt(glm::vec3(2.0f, 2.0f, 2.0f), glm::vec3(0.0f, 0.0f, 0.0f), glm::vec3(0.0f,
                                                                                                    0.0f, 1.0f));
        auto proj = glm::perspective(glm::radians(45.0f),
                                     context->getSwapChainExtent().width / (float)context->getSwapChainExtent().height, 0.1f, 10.0f);
        proj[1][1] *= -1;

        model = proj * view * model;

        memcpy(data, &model[0][0], size);
        return sizeof(model);
    }
private:
    VK_Context *context = nullptr;
    VK_Pipeline *pipeline = nullptr;
};

REGISTER_OBJECT(Node, cube)
